"""
优化工作流模块
提供经过优化的研究助理工作流实现
"""

from langgraph.graph import StateGraph, END
from src.research_core.state import AgentState
import logging

logger = logging.getLogger(__name__)

def create_optimized_workflow():
    """
    创建优化的工作流
    
    Returns:
        Pregel: 编译后的工作流
    """
    # 延迟导入节点函数，避免循环导入
    from src.research_core.nodes import (
        generate_search_query, 
        execute_web_search, 
        evaluate_results, 
        generate_final_answer
    )
    
    # 创建包装函数以适配LangGraph节点签名要求
    def generate_search_query_node(state):
        return generate_search_query(state)
    
    async def web_search_node(state):
        result = execute_web_search(state)
        # 如果是协程对象，则await它
        if hasattr(result, '__await__'):
            result = await result
        return result
    
    def evaluate_node(state):
        return evaluate_results(state)
    
    def generate_answer_node(state):
        return generate_final_answer(state)
    
    # 定义条件边函数
    def should_continue(state):
        """条件边函数：决定图是继续循环还是结束"""
        if state.get('research_complete', False):
            return "end"  # 去往终点
        else:
            return "continue"  # 回去重新生成查询

    # 构建图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("generate_query", generate_search_query_node)   # 添加生成查询节点
    workflow.add_node("web_search", web_search_node)          # 添加网络搜索节点
    workflow.add_node("evaluate", evaluate_node)              # 添加评估节点
    workflow.add_node("generate_answer", generate_answer_node)  # 添加生成答案节点

    # 设置入口点
    workflow.set_entry_point("generate_query")

    # 添加边
    workflow.add_edge("generate_query", "web_search")
    workflow.add_edge("web_search", "evaluate")

    # 添加条件边
    workflow.add_conditional_edges(
        "evaluate",
        should_continue,
        {
            "end": "generate_answer",  # 如果完成，去生成答案
            "continue": "generate_query"  # 如果继续，回到生成查询
        }
    )

    # 添加最后的边
    workflow.add_edge("generate_answer", END)

    return workflow

graph = create_optimized_workflow()
